"""
图像生成API路由
提供图像生成相关的API接口，包括任务创建、状态查询和结果获取
支持高并发场景和用户认证
"""

import logging
from fastapi import APIRouter, Depends, HTTPException, Request
from typing import Optional, Dict, Any
from pydantic import BaseModel, Field
from datetime import datetime
import uuid

from ..services.image_generation_service import image_generation_service
from ..services.character_service import load_character_profile
from ..services.jwt_service import jwt_required, send_token_request_background
from ..services.comfyui_service import comfyui_service
from ..core.rate_limiter import get_client_ip, check_rate_limit

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/image", tags=["image"])


class CreateImageRequest(BaseModel):
    """创建图像请求模型"""
    user_prompt: str = Field(..., min_length=1, max_length=200, description="用户提示词")
    ai_person_id: int = Field(..., description="关联的AI角色ID")
    reserve_id: Optional[str] = Field(None, description="预留ID")
    
    class Config:
        json_schema_extra = {
            "example": {
                "user_prompt": "一个穿着粉红色晚礼服的女孩在面包店拿着红豆包",
                "ai_person_id": 3,
                "reserve_id": "2abcd28a-b8d1-4ac9-8d59-821b5f9cde6c"
            }
        }


@router.post("/create", dependencies=[Depends(jwt_required)])
async def create_image(
    request: CreateImageRequest,
    req: Request
) -> Dict[str, Any]:
    """
    创建图像生成任务
    
    Args:
        request: 创建图像请求数据
        req: FastAPI请求对象
        
    Returns:
        Dict: 任务创建结果
    """
    try:
        # 速率限制检查
        client_ip = get_client_ip(req)
        try:
            check_rate_limit(client_ip)
        except Exception as e:
            raise HTTPException(status_code=429, detail="请求过于频繁，请稍后再试")
        
        # 获取用户信息
        user_data = getattr(req.state, 'user', None)
        if not user_data:
            raise HTTPException(status_code=401, detail="用户未认证")
        
        user_uuid = user_data.get('uuid')
        if not user_uuid:
            raise HTTPException(status_code=401, detail="无效的用户信息")
        
        # 获取角色配置信息
        try:
            character_profile = await load_character_profile(chat_id=request.ai_person_id)
            if not character_profile:
                logger.error(f"无法获取角色ID {request.ai_person_id} 的配置信息")
                send_token_request_background(request.reserve_id, False)
                raise HTTPException(status_code=404, detail="角色配置不存在")
            
            # 提取工作流类型和角色触发词
            work_flow = character_profile.get('work_flow')
            character_trigger = character_profile.get('persionImageDescription', '')
            
            if not work_flow:
                logger.error(f"角色配置缺少work_flow字段: {character_profile}")
                send_token_request_background(request.reserve_id, False)
                raise HTTPException(status_code=400, detail="角色配置缺少工作流信息")
            
            # 打印选中的工作流和角色描述信息
            logger.info(f"选中的工作流: {work_flow}")
            if character_trigger:
                logger.info(f"角色描述(persionImageDescription): {character_trigger}")
            else:
                logger.info("角色描述(persionImageDescription): 未设置，将使用默认配置")
            
        except HTTPException:
            raise
        except Exception as e:
            logger.error(f"获取角色配置失败: {e}", exc_info=True)
            send_token_request_background(request.reserve_id, False)
            raise HTTPException(status_code=500, detail="获取角色配置失败")
        
        # 获取客户端信息
        user_agent = req.headers.get("user-agent", "")
        
        # 调用业务服务创建任务，传入从角色配置获取的信息
        try:
            result = await image_generation_service.create_image_task(
                user_uuid=user_uuid,
                character_trigger=character_trigger,
                user_prompt=request.user_prompt,
                ai_person_id=request.ai_person_id,
                client_ip=client_ip,
                user_agent=user_agent,
                style_name=work_flow  # 传入工作流类型
            )
            
            # 处理业务逻辑返回的错误
            if "error" in result:
                error_code = result.get("code", 500)
                # 业务逻辑错误，发送取消请求
                send_token_request_background(request.reserve_id, False)
                raise HTTPException(status_code=error_code, detail=result["error"])
            
            # 图像任务创建成功，发送确认请求
            send_token_request_background(request.reserve_id, True)
            return result
        except HTTPException:
            # 重新抛出HTTP异常
            raise
        except Exception as e:
            # 其他异常，发送取消请求
            send_token_request_background(request.reserve_id, False)
            raise e
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"创建图像任务失败: {e}", exc_info=True)
        # 处理失败，发送取消请求
        send_token_request_background(request.reserve_id, False)
        raise HTTPException(status_code=500, detail="服务器内部错误")


@router.get("/task/{task_id}", dependencies=[Depends(jwt_required)])
async def get_task_status(
    task_id: str,
    req: Request
) -> Dict[str, Any]:
    """
    获取任务状态
    
    Args:
        task_id: 任务ID
        req: FastAPI请求对象
        
    Returns:
        Dict: 任务状态信息
    """
    try:
        # 验证任务ID格式
        try:
            uuid.UUID(task_id)
        except ValueError:
            raise HTTPException(status_code=400, detail="无效的任务ID格式")
        
        # 获取用户信息
        user_data = getattr(req.state, 'user', None)
        if not user_data:
            raise HTTPException(status_code=401, detail="用户未认证")
        
        user_uuid = user_data.get('uuid')
        if not user_uuid:
            raise HTTPException(status_code=401, detail="无效的用户信息")
        
        # 获取任务状态
        result = await image_generation_service.get_task_status(task_id, user_uuid)
        
        # 处理业务逻辑返回的错误
        if "error" in result:
            error_code = result.get("code", 500)
            raise HTTPException(status_code=error_code, detail=result["error"])
        
        return result
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"获取任务状态失败: {task_id} - {e}", exc_info=True)
        raise HTTPException(status_code=500, detail="服务器内部错误")


@router.get("/servers/status", dependencies=[Depends(jwt_required)])
async def get_servers_status(req: Request) -> Dict[str, Any]:
    """
    获取ComfyUI服务器状态
    
    Args:
        req: FastAPI请求对象
        
    Returns:
        Dict: 服务器状态信息
    """
    try:
        # 获取用户信息验证权限
        user_data = getattr(req.state, 'user', None)
        if not user_data:
            raise HTTPException(status_code=401, detail="用户未认证")
        
        # 获取服务器健康状态
        health_status = await comfyui_service.check_server_health()
        
        # 获取负载均衡器统计信息
        load_balancer_stats = await comfyui_service.get_load_balancer_stats()
        
        return {
            "server_health": health_status,
            "load_balancer": load_balancer_stats,
            "timestamp": datetime.utcnow().isoformat()
        }
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"获取服务器状态失败: {e}", exc_info=True)
        raise HTTPException(status_code=500, detail="服务器内部错误") 